from sklearn.model_selection import StratifiedShuffleSplit
import pandas as pd
import numpy as np
wj='lenses'
cs=[0.9,0.7,0.5]



def stratifid(df, target, test_sz, fileFolder): # 带默认值的形参必须放在最后
    i = 1
    xl=100-test_sz*100
    split = StratifiedShuffleSplit(n_splits=10, test_size=test_sz, random_state=0)
    for tr_idx, te_idx in split.split(df, df[target]):
        train = df.loc[tr_idx]
        test = df.loc[te_idx]
        train.to_csv(('new\\%s\\%i%%\\train\{}train.txt'%(wj,xl)).format(str(i)), sep='\t', header=0, index=0)
        test.to_csv(('new\\%s\\%i%%\\test\{}test.txt'%(wj,xl)).format(str(i)), sep='\t', header=0, index=0)
        print("第%d个数据集处理完毕" % i)
        i += 1
        # print_class_label_split(train, test) # 检查标签分布
    # print("全部数据处理完毕")
    # return train, test


#定义get_class_distribution函数，它采用y的类别标签作为参数，返回一个字典，键是类别标签，值是这类记录数占总数的百分比分布。返回类别标签的分布情况。在随后的函数里调用这个函数，就可以了解训练集和测试集里的类别分布。
def get_class_distribution(y):
    d = {}
    set_y = set(y)
    # print(set_y)
    for y_label in set_y:
        no_elements = len(np.where(y == y_label)[0])
        d[y_label] = no_elements
    dist_percentage = {class_label: count/(1.0*sum(d.values())) for class_label, count in d.items()}
    return dist_percentage


#定义print_class_label_split函数，把训练集和测试集作为参数。
def print_class_label_split(train, test):
    #打印训练集类别分布
    y_train = train.loc[:,0]
    train_distribution = get_class_distribution(y_train)
    print('\n Train data set class label distribution')
    for k, v in train_distribution.items():
        print('Class label = %d, percentage records = %.2f)'%(k, v))

    #打印测试集类别分布
    y_test = test.loc[:,0]
    test_distribution = get_class_distribution(y_test)
    print('\n Test data set class label distribution')
    for k, v in test_distribution.items():
        print('Class label = %d, percentage records = %.2f)'%(k, v))


if __name__ == '__main__':
    fileName = "0213/%s.txt"%wj
    fileFolder = fileName[:5]
    # fileFolder = "awa_eigen"
    file = pd.read_csv(fileName, encoding='utf-8', sep='\s+', header=None, engine='python') # 如果数据间是标准化格式用“\t”,否则用"\s+"
    file = file.sample(frac=1.0)
    for i in cs:
        stratifid(file, [0],i , fileFolder)

